﻿/*
 * FfmpegMediaReader.cpp
 *
 *  Created on: 2017年3月15日
 *      Author: chuanjiang.zh
 */

#include "FfmpegMediaReader.h"
#include "CLog.h"
#include "DateTime.h"
#include "TStringCast.h"
#include "TStringUtil.h"
#include "FfmpegUtil.h"
#include "Path.h"
#include "H264PropParser.h"
#include "H265PropParser.h"
#include "Socket.h"
#include "TSafeStr.h"

#include "libavutil/intreadwrite.h"

using namespace av;



static int check_interrupt(void* opaque)
{
    FfmpegMediaReader* obj = (FfmpegMediaReader*)opaque;
    return obj->checkInterrupt();
}

FfmpegMediaReader::FfmpegMediaReader():
		m_dict(),
		m_fmtContext(),
		m_videoIdx(-1),
		m_audioIdx(-1),
		m_hasBFrame(),
		m_state(UMEDIA_STATE_STOPPED),
        m_interrupted(),
        m_paused(),
		m_isFile(),
		m_curTime(),
		m_transformNalu(),
        m_bsfc(),
        m_audioTsCustomed(),
        m_audioTs()
{
}

FfmpegMediaReader::~FfmpegMediaReader()
{
	close();
}

int FfmpegMediaReader::open(const char* url, const char* params)
{
    CLog::info("FfmpegMediaReader(%p) open begin. %s\n", this, url);

	if (isOpen())
	{
		close();
	}

	m_url = url;
	m_params = params;
	m_isFile = isFile();
	
	int rc = openInput(url, params);
	if (rc != 0)
	{
		CLog::warning("failed to open url. rc:%d, url: %s\n", rc, url);
		return rc;
	}

	std::string fname(m_fmtContext->iformat->name);
	m_fmtName = comn::StringUtil::getHead(fname, ',');

	if (comn::StringUtil::startsWith(url, "rtmp://"))
	{
		m_transformNalu = true;
	}
	else if (fname.find("mp4") != std::string::npos)
	{
		m_transformNalu = true;
	}
	else if (fname.find("flv") != std::string::npos)
	{
		m_transformNalu = true;
	}
	else if (fname.find("matroska,webm") != std::string::npos)
	{
		m_transformNalu = true;
	}

	fetchMediaFormat();

    onOpen(m_fmtContext);

    //CLog::info("FfmpegMediaReader(%p) open end. audio codec:%d, channels:%d, rate:%d\n",
    //    this, m_format.m_audioCodec, m_format.m_channels, m_format.m_sampleRate);

	return 0;
}

void FfmpegMediaReader::close()
{
    if (m_interrupted > 0)
    {
        interrupt();
    }

	closeInput();

	closeFilter();

    m_paused = false;
    m_isFile = false;
    m_curTime = 0;
    m_transformNalu = false;
    m_audioTs = 0;
}

bool FfmpegMediaReader::isOpen()
{
	return isInputOpen();
}

bool FfmpegMediaReader::getFormat(UMediaFormat& fmt)
{
	fmt = m_format;
	return true;
}

int FfmpegMediaReader::getDuration()
{
	return m_format.m_duration;
}

bool FfmpegMediaReader::play()
{
	if (!m_fmtContext)
	{
		return false;
	}

    if (!m_paused)
    {
        return false;
    }

    int ret = av_read_play(m_fmtContext);
    return ret == 0;
}

bool FfmpegMediaReader::pause()
{
	if (!m_fmtContext)
	{
		return false;
	}

    m_paused = true;
	int ret = av_read_pause(m_fmtContext);
    return ret == 0;
}

void FfmpegMediaReader::stop()
{
	//
    m_paused = true;
}

int FfmpegMediaReader::getState()
{
	comn::AutoCritSec lock(m_cs);
	return m_state;
}

bool FfmpegMediaReader::seekable()
{
	return (m_format.m_duration > 0);
}

bool FfmpegMediaReader::seek(int64_t offset)
{
	if (!m_fmtContext)
	{
		return false;
	}

	int64_t pos = offset;
	int ret = av_seek_frame(m_fmtContext, -1, pos, 0);
    return ret == 0;
}

int64_t FfmpegMediaReader::getTime()
{
	return m_curTime;
}

bool FfmpegMediaReader::setScale(float scale)
{
	return false;
}

float FfmpegMediaReader::getScale()
{
	return 1.0;
}

//int FfmpegMediaReader::read(UMediaPacket& pkt)
//{
//	return read(*pkt.get());
//}

int FfmpegMediaReader::read(AVPacket& pkt)
{
	if (!m_fmtContext)
	{
		return ENOENT;
	}

    int rc = doRead(pkt);

	for (int i = 0; i < 5 && (rc == AVERROR(EAGAIN)); ++ i)
	{
		rc = doRead(pkt);
	}

    if (rc == 0)
    {
        if (pkt.stream_index == UMEDIA_TYPE_VIDEO)
        {
            if (!convert(pkt))
            {
                return EAGAIN;
            }
        }
        else if (pkt.stream_index == UMEDIA_TYPE_AUDIO)
        {
            if (m_audioTsCustomed)
            {
                customAudioTs(pkt);
            }
        }
    }

    return rc;
}


int FfmpegMediaReader::openInput(const std::string& url, const std::string& params)
{
    comn::DateTime timeStart = comn::DateTime::now();

	parseParams(&m_dict, params);

	AVInputFormat* fmt = NULL;
    m_fmtContext = avformat_alloc_context();

    m_fmtContext->interrupt_callback.callback = check_interrupt;
    m_fmtContext->interrupt_callback.opaque = this;

    ffmpeg::av_dict_get_int(m_dict, "max_delay", m_fmtContext->max_delay);

	int rc = avformat_open_input(&m_fmtContext, url.c_str(), fmt, &m_dict);
	if (rc != 0)
	{
		return ENOSYS;
	}

	 //= AV_TIME_BASE / 4;
	ffmpeg::av_dict_get_int(m_dict, "max_analyze_duration", m_fmtContext->max_analyze_duration);

	ffmpeg::av_dict_get_int(m_dict, "probesize", m_fmtContext->probesize);

	rc = avformat_find_stream_info(m_fmtContext, NULL);
	if (rc != 0)
	{
		avformat_close_input(&m_fmtContext);
		return ENOSPC;
	}

	comn::DateTime timeNow = comn::DateTime::now();
	time_t ms = timeNow.subtract(timeStart);
	//CLog::info("elapse of open url. %d\n", ms);

	AVCodec* vCodec = NULL;
	AVCodec* aCodec = NULL;
	m_videoIdx = av_find_best_stream(m_fmtContext, AVMEDIA_TYPE_VIDEO, -1, -1, &vCodec, 0);
	m_audioIdx = av_find_best_stream(m_fmtContext, AVMEDIA_TYPE_AUDIO, -1, m_videoIdx, &aCodec, 0);

	return 0;
}

void FfmpegMediaReader::closeInput()
{
	if (m_fmtContext)
	{
		avformat_close_input(&m_fmtContext);
	}

    if (m_bsfc)
    {
        av_bitstream_filter_close(m_bsfc);
        m_bsfc = NULL;
    }

	m_videoIdx = -1;
	m_audioIdx = -1;

	if (m_dict)
	{
		av_dict_free(&m_dict);
	}
}

bool FfmpegMediaReader::isInputOpen()
{
	return (m_fmtContext != NULL);
}

void FfmpegMediaReader::fetchMediaFormat()
{
	if (m_videoIdx >= 0)
	{
		AVStream* stream = m_fmtContext->streams[m_videoIdx];
        AVCodecContext* codec = stream->codec;
		m_format.m_codec = stream->codec->codec_id;
        m_format.m_profile = codec->profile;
		m_format.m_width = stream->codec->width;
		m_format.m_height = stream->codec->height;
		m_format.m_bitrate = stream->codec->bit_rate;

		if (stream->avg_frame_rate.den > 0)
		{
			m_format.m_framerate = ceil(av_q2d(stream->avg_frame_rate));
		}

        m_format.m_clockRate = AV_TIME_BASE;

        if (codec->extradata_size > 0)
        {
            if (m_format.m_codec == UMEDIA_CODEC_H264)
            {
                std::string sps;
                std::string pps;
				if (av::H264PropParser::splitPropSet(codec->extradata, codec->extradata_size, sps, pps))
				{

					std::string vprop;
					vprop.append((const char*)av::H264PropParser::getStartCode(),
						av::H264PropParser::START_CODE_LENGTH);
					vprop.append(sps);
					vprop.append((const char*)av::H264PropParser::getStartCode(),
						av::H264PropParser::START_CODE_LENGTH);
					vprop.append(pps);

					m_vprop = vprop;
				}
            }
			else if (m_format.m_codec == UMEDIA_CODEC_HEVC)
			{
				std::string vps;
				std::string sps;
				std::string pps;
				if (av::H265PropParser::splitPropSet(codec->extradata, codec->extradata_size, vps, sps, pps))
				{
					std::string vprop;
					vprop.append((const char*)av::H264PropParser::getStartCode(), av::H264PropParser::START_CODE_LENGTH);
					vprop.append(vps);
					vprop.append((const char*)av::H264PropParser::getStartCode(), av::H264PropParser::START_CODE_LENGTH);
					vprop.append(sps);
					vprop.append((const char*)av::H264PropParser::getStartCode(), av::H264PropParser::START_CODE_LENGTH);
					vprop.append(pps);

					m_vprop = vprop;
				}
			}

			if (needTransformNalu())
			{
				openFilter();

				m_format.setVideoProp(m_vprop.c_str(), m_vprop.size());
			}
			else
			{
				m_format.setVideoProp(codec->extradata, codec->extradata_size);
			}
        }

		m_hasBFrame = stream->codec->has_b_frames;

		if (stream->duration != AV_NOPTS_VALUE)
		{
			m_format.m_videoDuration = stream->duration * 1000 / AV_TIME_BASE;
		}
		
	}

	if (m_audioIdx >= 0)
	{
		AVStream* stream = m_fmtContext->streams[m_audioIdx];
        AVCodecContext* codec = stream->codec;
		m_format.m_audioCodec = stream->codec->codec_id;
        m_format.m_audioProfile = codec->profile;

		m_format.m_channels = stream->codec->channels;
		m_format.m_sampleRate = stream->codec->sample_rate;
		m_format.m_sampleBits = stream->codec->bits_per_coded_sample;
		m_format.m_frameSize = stream->codec->frame_size;

        m_format.m_audioRate = AV_TIME_BASE;

        if (codec->extradata_size > 0)
        {
            m_format.setAudioConfig(codec->extradata, codec->extradata_size);
        }
	}

    if (m_fmtContext->duration != AV_NOPTS_VALUE)
    {
	    m_format.m_duration = m_fmtContext->duration * 1000 / AV_TIME_BASE;
    }

}

void FfmpegMediaReader::setState(UMediaState state)
{
	comn::AutoCritSec lock(m_cs);
	m_state = state;
}

int FfmpegMediaReader::parseParams(AVDictionary **pm, const std::string& params)
{
	int rc = av_dict_parse_string(pm, params.c_str(), "=", ";", 0);
	//av_dict_set(pm, "stimeout", "5000000", 0);

	av_dict_set(pm, "connect", "1", 0);

	return rc;
}

void FfmpegMediaReader::interrupt()
{
    comn::AutoCritSec lock(m_cs);
    m_interrupted = true;
}

bool FfmpegMediaReader::isLive()
{
	if (isFile())
	{
		return false;
	}

	return (m_format.m_duration <= 0);
}

const char* FfmpegMediaReader::getFormatName()
{
	return m_fmtName.c_str();
}

size_t FfmpegMediaReader::getMetaData(const char* name, char* value, size_t size)
{
	std::string key(name);
	if (key == "bframe")
	{
		std::string str = comn::StringCast::toString(m_hasBFrame);
		return comn::copyStr(value, size, str);
	}

	return 0;
}

int FfmpegMediaReader::checkInterrupt()
{
    comn::AutoCritSec lock(m_cs);
    return m_interrupted;
}

int FfmpegMediaReader::doRead(AVPacket& pkt)
{
	comn::DateTime timeStart = comn::DateTime::now();

	//av_init_packet(&pkt);
	int rc = av_read_frame(m_fmtContext, &pkt);
	if (rc == 0)
	{
		AVStream* stream = m_fmtContext->streams[pkt.stream_index];
        if (stream->start_time != AV_NOPTS_VALUE)
        {
            pkt.pts -= stream->start_time;
            pkt.dts -= stream->start_time;
        }
		else
		{
		}

        if (stream->time_base.den > 0)
        {
            AVRational tb_dst = av_make_q(1, AV_TIME_BASE);
            av_packet_rescale_ts(&pkt, stream->time_base, tb_dst);
        }

		if (pkt.stream_index == m_videoIdx)
		{
			pkt.stream_index = UMEDIA_TYPE_VIDEO;

			if ((pkt.dts == AV_NOPTS_VALUE) && isFile())
			{
				pkt.dts = m_curTime + pkt.duration;
				pkt.pts = pkt.dts;
			}

			m_curTime = pkt.dts;
		}
		else if (pkt.stream_index == m_audioIdx)
		{
			pkt.stream_index = UMEDIA_TYPE_AUDIO;

			m_curTime = pkt.dts;
		}
		else
		{
			av_packet_unref(&pkt);
			rc = AVERROR(EAGAIN);
		}

        time_t ms = timeStart.getElapse();
		if (ms > 1000)
		{
			CLog::info("-- elapse of doRead. %d. url: %s\n", (int)ms, m_url.c_str());
		}

		//av_packet_unref(&pkt);
		//rc = AVERROR(EAGAIN);
		//return rc;
	}

	return rc;
}

bool FfmpegMediaReader::isFile()
{
	return comn::Path::exist(m_url);
}

bool FfmpegMediaReader::needTransformNalu()
{
	return (m_format.m_codec == UMEDIA_CODEC_H264 || m_format.m_codec == UMEDIA_CODEC_HEVC) && m_transformNalu;
}

void FfmpegMediaReader::openFilter()
{
}

void FfmpegMediaReader::closeFilter()
{
}

void FfmpegMediaReader::onOpen(AVFormatContext* fmtContext)
{
    if (m_videoIdx >= 0)
    {
        AVStream* stream = fmtContext->streams[m_videoIdx];
		if (m_transformNalu)
		{
			if (stream->codec->codec_id == AV_CODEC_ID_H264)
			{
				m_bsfc = av_bitstream_filter_init("h264_mp4toannexb");
			}
			else if (stream->codec->codec_id == AV_CODEC_ID_HEVC)
			{
				m_bsfc = av_bitstream_filter_init("hevc_mp4toannexb");
			}
		}
    }
}

bool FfmpegMediaReader::convert(AVPacket& pktIn)
{
    if (m_bsfc)
    {
        uint8_t* outBuf = NULL;
        int outSize = 0;
        AVStream* stream = m_fmtContext->streams[m_videoIdx];
        int ret = av_bitstream_filter_filter(m_bsfc, stream->codec, NULL, &outBuf, &outSize, pktIn.data, pktIn.size, 0);
		if (ret >= 0)
		{
			AVPacket pktOut;
			av_init_packet(&pktOut);
			av_packet_copy_props(&pktOut, &pktIn);

			av_packet_unref(&pktIn);
			av_init_packet(&pktIn);
			av_packet_from_data(&pktIn, outBuf, outSize);
			av_packet_copy_props(&pktIn, &pktOut);
		}
    }
    else if ((pktIn.flags & AV_PKT_FLAG_KEY) != 0)
    {
        H264PropParser::NaluPacket nalu;
        H264PropParser::parseNalu(pktIn.data, pktIn.size, nalu);
        if (nalu.type == H264PropParser::NaluPacket::NALU_IFRAME && m_vprop.size() > 0)
        {
            AVPacket pktOut;
            av_init_packet(&pktOut);
            av_new_packet(&pktOut, (int)(m_vprop.size() + pktIn.size));
            av_packet_copy_props(&pktOut, &pktIn);
            
            memcpy(pktOut.data, m_vprop.c_str(), m_vprop.size());
            memcpy(pktOut.data + m_vprop.size(), pktIn.data, pktIn.size);

            av_packet_unref(&pktIn);
            av_init_packet(&pktIn);
            
            av_packet_ref(&pktIn, &pktOut);
            av_packet_copy_props(&pktIn, &pktOut);

            av_packet_unref(&pktOut);
        }
    }
    return true;
}

void FfmpegMediaReader::customAudioTs(bool enabled)
{
    m_audioTsCustomed = enabled;
}

void FfmpegMediaReader::customAudioTs(AVPacket& pkt)
{
    //int64_t delta = (int)(m_audioTs - pkt.dts)/1000;
    //CLog::debug("customAudioTs. dts: %lld, pts: %lld, duration: %d, audioTs: %lld, delta: %lld ms\n",
    //    pkt.dts/1000, pkt.pts / 1000, (int)pkt.duration / 1000, m_audioTs / 1000, delta);

    pkt.dts = m_audioTs;
    pkt.pts = pkt.dts;

    m_audioTs += pkt.duration;
}

int FfmpegMediaReader::getBFrame()
{
	return m_hasBFrame;
}

double FfmpegMediaReader::countFrameRate(int* frameCount, int64_t* frameDuration)
{
	if (m_format.m_codec == 0)
	{
		return 0;
	}

	int maxCount = (frameCount != nullptr) ? *frameCount : 60;

	int64_t duration = 0;
	int64_t startTs = 0;
	int64_t stopTs = 0;
	int64_t pktCount = 0;

	while ((maxCount <= 0) || (maxCount > 0 && pktCount < maxCount))
	{
		AVPacket pkt;
		av_init_packet(&pkt);

		int ret = doRead(pkt);
		if (ret != 0)
		{
			break;
		}

		if (pkt.stream_index == m_videoIdx)
		{
			if (pktCount == 0)
			{
				startTs = pkt.dts;
				stopTs = pkt.dts;
			}
			else
			{
				stopTs = pkt.dts;
			}

			pktCount ++;
		}

		av_packet_unref(&pkt);
	}

	duration = (stopTs - startTs);
	if (duration <= 0)
	{
		return 0;
	}

	double rate = ((double)(pktCount * AV_TIME_BASE)) / duration;

	if (frameCount)
	{
		*frameCount = pktCount;
	}

	if (frameDuration)
	{
		*frameDuration = duration;
	}

	return rate;
}